import torch
import torch.nn as nn


class MultiScaleCNN(nn.Module):
    def __init__(self, num_classes=10,
                 use_branch1=True, use_branch2=True, use_branch3=True):
        super(MultiScaleCNN, self).__init__()

        self.use_branch1 = use_branch1  # 3×3分支
        self.use_branch2 = use_branch2  # 5×5分支
        self.use_branch3 = use_branch3  # 7×7分支

        # 分支1：3×3卷积
        if use_branch1:
            self.branch1 = nn.Sequential(
                nn.Conv2d(1, 32, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.MaxPool2d(2, 2),
                nn.Conv2d(32, 64, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.MaxPool2d(2, 2),
                nn.Flatten()
            )

        # 分支2：5×5卷积
        if use_branch2:
            self.branch2 = nn.Sequential(
                nn.Conv2d(1, 32, kernel_size=5, padding=2),
                nn.ReLU(),
                nn.MaxPool2d(2, 2),
                nn.Conv2d(32, 64, kernel_size=5, padding=2),
                nn.ReLU(),
                nn.MaxPool2d(2, 2),
                nn.Flatten()
            )

        # 分支3：7×7卷积
        if use_branch3:
            self.branch3 = nn.Sequential(
                nn.Conv2d(1, 32, kernel_size=7, padding=3),
                nn.ReLU(),
                nn.MaxPool2d(2, 2),
                nn.Conv2d(32, 64, kernel_size=7, padding=3),
                nn.ReLU(),
                nn.MaxPool2d(2, 2),
                nn.Flatten()
            )

        # 动态计算全连接层输入维度
        fc_input_dim = 0
        if use_branch1:
            fc_input_dim += 64 * 7 * 7
        if use_branch2:
            fc_input_dim += 64 * 7 * 7
        if use_branch3:
            fc_input_dim += 64 * 7 * 7

        self.fc = nn.Sequential(
            nn.Linear(fc_input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        outputs = []
        if self.use_branch1:
            outputs.append(self.branch1(x))
        if self.use_branch2:
            outputs.append(self.branch2(x))
        if self.use_branch3:
            outputs.append(self.branch3(x))

        fused = torch.cat(outputs, dim=1)
        return self.fc(fused)